import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from PyEMD import CEEMDAN

def fetch_futures_data(ticker, start_date, end_date):
    """从本地CSV文件获取期货数据"""
    # 根据ticker选择对应的本地文件
    file_map = {
        'GC=F': 'data/SA0SH.csv',  # 黄金期货
        'SI=F': 'data/SF0SH.csv'   # 白银期货
    }
    
    # 读取数据
    data = pd.read_csv(file_map[ticker], index_col='date_time')
    
    # 仅保留收盘价
    data = data[['close']]
    data.columns = ['Close']  # 保持列名一致
    
    # 按日期范围筛选
    data = data.loc[start_date:end_date]
    
    return data

def preprocess_data(data):
    """数据预处理：去除NaN值，去趋势"""
    # 填补缺失值
    data = data.fillna(method='ffill').fillna(method='bfill')
    
    # 去趋势（可选）
    # data = data - data.rolling(window=20).mean()
    
    return data.values.flatten()  # 转换为1D数组

def ceemdan_decomposition(data, num_imfs=3, num_ensembles=100, noise_strength=0.1):
    """使用CEEMDAN进行分解"""
    ceemdan = CEEMDAN(range_thr=3.0, num_ensembles=num_ensembles, noise_strength=noise_strength)
    imfs = ceemdan(data)
    
    return imfs

def plot_decomposition(data, imfs, title):
    """绘制分解结果"""
    plt.figure(figsize=(12, 8))
    
    # 绘制原始数据
    plt.subplot(len(imfs) + 1, 1, 1)
    plt.plot(data)
    plt.title(f'Original Data - {title}')
    plt.grid(alpha=0.3)
    plt.tight_layout()
    
    # 绘制每个IMF
    for i, imf in enumerate(imfs):
        plt.subplot(len(imfs) + 1, 1, i + 2)
        plt.plot(imf)
        plt.title(f'IMF {i+1}')
        plt.grid(alpha=0.3)
        
    plt.tight_layout()
    plt.show()

def main():
    # 参数设置
    ticker = 'GC=F'  # 黄金期货代码
    start_date = '2010-01-01'
    end_date = '2025-12-31'
    
    # 获取数据
    print("获取数据...")
    data = fetch_futures_data(ticker, start_date, end_date)
    
    # 数据预处理
    print("数据预处理...")
    processed_data = preprocess_data(data)
    
    # CEEMDAN分解
    print("CEEMDAN分解...")
    imfs = ceemdan_decomposition(processed_data, num_imfs=3)
    
    # 可视化结果
    print("绘制结果...")
    plot_decomposition(processed_data, imfs, ticker)

if __name__ == "__main__":
    main()
